Flask のテストクライアントをカスタムしてデフォルトのHTTPヘッダを設定する

Pythonの軽量フレームワークFlaskで自動テストをするためのテストクライアントをカスタムして、リクエスト毎にデフォルトのHTTPヘッダを設定する方法をご紹介します

#Python

#Flask

西田将幸

2018.08.22

この記事は公開されてから1年以上経過しています。情報が古い可能性がありますので、ご注意ください。

Flask のテストクライアントをカスタムしてデフォルトのHTTPヘッダを設定する

西田@大阪です

Pythonの軽量フレームワークFlask では自動テストをするためのテストクライアントが用意されています

今回はそのテストクライアントをカスタムして、リクエスト毎にデフォルトのHTTPヘッダを設定する方法をご紹介します

試した環境

Python 3.6.2
Flask 1.0.2

やりたいこと

今回はREST APIのテストをします

登場するAPIは以下の２つです

/authenticate 認証を行いaccess_tokenを返します
/products 商品情報を返します。認証が必要です

/authenticate のレスポンスに含まれる access_token を Authorizationヘッダに含めると/productsにアクセスすることができます。

/auhenticate

リクエスト

{
    "username": "user01",
    "password": "xxxxxx"
}

レスポンス

{
    "access_token": "xxx.xxxx.xxx"
}

/products

Authorization HTTPヘッダに/authenticateのレスポンスに含まれるaccess_tokenを設定することで認証できます

ヘッダ

Authorization: Bearer xxx.xxx.xxx

レスポンス

[
    {
        "id": 123,
        "name": "サンプル商品01"
    },
    {
        "id": 124,
        "name": "サンプル商品02"
    }
]

実装

カスタムしたtest_clientのクラスを作成します

FlaskClient を継承して作成します

from flask.testing import FlaskClient

class SecuredFlaskClient(FlaskClient):
    def __init__(self, *args, **kwargs):
        # 拡張したパラメーターをそのまま FlaskClient の
        # コンストラクタに渡さないために pop でとりだします
        access_token = kwargs.pop('access_token')
        super(SecuredFlaskClient, self).__init__(*args, **kwargs)

        # FlaskClient.envron_baseに'HTTP_' から始まるキーで値を設定すると、
        # Flask(WSGI)に渡すデフォルトのリクエストヘッダーになることを利用して
        # 初期化時に渡された access_token をリクエスト毎に 
        # Authorization リクエストヘッダに追加します
        self.environ_base['HTTP_AUTHORIZATION'] = f"Bearer {access_token}"

test_client を返すメソッドを用意します

from flaskr import flaskr

# デフォルトのテストクライアントを返します。認証のいらないAPIをテストするのに使います
def test_client():
    # test_client のクラスをデフォルトのFlaskClientに設定する
    flaskr.app.test_client_class = None
    return flaskr.app.test_client()

# カスタムされたテストクライアントを返します。認証の必要なAPIをテストするのに使います
def secured_test_client(username, password):
    # 認証を行うAPIを呼び出します
    rv = test_client().post('/authenticate', json=dict(username=username, password=password))

    # test_client のクラスを作成した SecuredFlaskClient に設定する
    flaskr.app.test_client_class = SecuredFlaskClient

    # テストクライアントに認証APIのレスポンスの `access_token`　を
    # 渡してインスタンスを生成します
    return flaskr.app.test_client(access_token=rv.json['access_token'])